import re

import pymysql

from llm.DeepSeekLLM import DeepSeekLLM
from util.T2llm import T2llm
from llm.prompt.native_to_sql_prompt import _DEFAULT_TEMPLATE as native_to_sql_template
from llm.prompt.native_to_data_analysis_prompt import _DEFAULT_TEMPLATE as data_analysis_prompt_template



class DButil:

    def __init__(self, host='127.0.0.1', user='root', password='tangzhangss', database='test'):
        self.host = host
        self.user = user
        self.password = password
        self.database = database

    def _get_conn(self):
        conn = pymysql.connect(host=self.host,  user = self.user, password=self.password, database=self.database)
        return  conn

    """
    :return  [(表名,建表信息) ]
    """
    def get_all_table_create_info(self)-> list[tuple]:
        conn = self._get_conn()
        cursor = conn.cursor()
        # 查询所有表
        # 获取所有表名
        cursor.execute("SHOW TABLES")
        tables = cursor.fetchall()
        tables_info=[]
        for table in tables:
            cursor.execute("SHOW CREATE TABLE " + table[0])
            info = cursor.fetchone()
            tables_info.append(info)
            # print(info)
        # 关闭连接
        cursor.close()
        conn.close()
        return tables_info

    """
    自然语言转SQL
    :param query 自然语言
    :param table_threshold table的获取上线数
    :return sql语句 如果没有生成SQL语句则返回None
    """
    def natural_language_to_sql(self, query, table_threshold=None):

        table_create_infos=self._get_table_create_info(query, table_threshold)
        #构建对话信息
        prompt = native_to_sql_template.replace('{db_name}', self.database).replace('{table_info}', '\n\n'.join([item[0] for item in table_create_infos])).replace('{input}', query)
        llm = DeepSeekLLM()
        result = llm.get(prompt)
        # 结果提取
        # 使用 re.search 查找匹配项
        match = re.search(r'```sql(.*?)```', result, re.DOTALL)
        if match:
            result = match.group(1).strip()
        else:
            result = None
        return  result

    """
        自然语言转分析图表
        :param query 自然语言
        :param table_threshold table的获取上线数
        :return 数据分析json数据 
    """
    def natural_language_to_data_analysis(self, query, table_threshold=None):

        table_create_infos = self._get_table_create_info(query, table_threshold)
        # 构建对话信息
        prompt = data_analysis_prompt_template.replace('{table_info}', '\n\n'.join(
            [item[0] for item in table_create_infos])).replace('{input}', query)
        llm = DeepSeekLLM()
        result = llm.get(prompt)

        # 结果提取
        # 使用 re.search 查找匹配项
        match = re.search(r'```json(.*?)```', result, re.DOTALL)
        if match:
            result = match.group(1).strip()
        else:
            result = None

        return result

    """
      获取table的结构信息 create
    """
    def _get_table_create_info(self, query, table_threshold):
        table_infos = self.get_all_table_create_info()
        t2llm = T2llm()
        # 将query 和 table_infos都向量化 然后计算相似度
        scores = t2llm.calculate_cosine(query, [item[1] for item in table_infos])
        # 获取建表语句
        table_create_infos = []
        for tableinfo, score in zip(table_infos, scores):
            table_create_infos.append((tableinfo[1], score))
        # 按照相似度排序
        # 将所有符合要求的文档按分数从高到低排序。
        table_create_infos.sort(key=lambda x: x[1], reverse=True)

        table_create_infos = table_create_infos[:table_threshold] if table_threshold else table_create_infos
        return table_create_infos

if __name__ == '__main__':
    dbutil =  DButil()
    # table_info = dbutil.get_all_table_create_info()
    # print(table_info)
    # sql = dbutil.natural_language_to_sql("获取2017年入职的人数")
    sql = dbutil.natural_language_to_data_analysis("按照入职时间统计人员信息")
    print(sql)
    #
